1457E - New Game Plus - CodeForces Solution

#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <vector>
#include <algorithm>
#include <cmath>
#include <set>
#include <climits>
#include <queue>
#include <numeric>

using namespace std;
 
#define ll long long
#define pb emplace_back
#define xx first
#define yy second

void solve() {
	ll n, k;
	cin >> n >> k;
	vector<ll> vec(n, 0);
	for(int i = 0; i < n; i++) {
		cin >> vec[i];
	}

	sort(vec.begin(), vec.end(), greater<ll>());

	ll segs = k + 1;
	vector<ll> sum(segs);
	vector<ll> bonus(segs);

	int i = 0;
	for (; i < n; ++i) {
		sum[0] += bonus[0];
		bonus[0] += vec[i];
		if (bonus[0] < 0) break;
	}

	i++;

	int j = 1;
	for (; i < n; ++i) {
		sum[j % segs] += bonus[j % segs];
		bonus[j % segs] += vec[i];
		++j;
	}

	ll ret = accumulate(sum.begin(), sum.end(), 0ll);
	cout << ret << endl;

}

int main() {
	solve();
}